package com.hdd.test;

import java.util.concurrent.TimeUnit;


/**
 * @BelongsProject: demo2
 * @BelongsPackage: com.hdd.test
 * @Author: songYuHang
 * @CreateTime: 2025-04-19  11:43
 * @Description: TODO
 * @Version: 1.0
 */
public class TokenBucket {

    // 1.令牌桶容量
    private int capacity;
    // 2.当前令牌数量
    private double tokens;
    // 3.令牌桶的填充速率
    private double rate;
    // 4.上次填充令牌的时间(毫秒)
    private long lastFillTime;

    public TokenBucket(int capacity, double rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = 0;
        this.lastFillTime = System.currentTimeMillis();
    }

    /**
     * 尝试消耗指定数量的令牌
     * @param tokensToConsume 要消耗的令牌数量
     * @return 如果有足够的令牌则返回 true，否则返回 false
     */
    public synchronized boolean consume(int tokensToConsume) {
        refill();
        if (tokens >= tokensToConsume) {
            tokens -= tokensToConsume;
            return true;
        }
        return false;
    }

    /**
     * 根据时间差更新令牌数量
     */
    private void refill() {
        long now = System.currentTimeMillis();
        // 计算从上次更新到现在应该生成的令牌数量
        double newTokens = (now - lastFillTime) * rate / TimeUnit.SECONDS.toMillis(1);
        tokens = Math.min(capacity, tokens + newTokens);
        lastFillTime = now;
    }

    public static void main(String[] args) {
        // 创建一个容量为 100，每秒生成 10 个令牌的令牌桶
        TokenBucket bucket = new TokenBucket(100, 10);

        // 模拟请求
        for (int i = 0; i < 20; i++) {
            if (bucket.consume(1)) {
                System.out.println("Request " + i + " is allowed.");
            } else {
                System.out.println("Request " + i + " is rate-limited.");
            }
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

}
